/******************************************************************************
 * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
 * libkperf licensed under the Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *     http://license.coscl.org.cn/MulanPSL2
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR
 * PURPOSE.
 * See the Mulan PSL v2 for more details.
 * Author:
 * Create: 2025-05-13
 * Description: Collection capability for LLC miss ratio and thread binding
 * Current capability: Bind threads to even/odd CPU cores based on LLC miss ratio
 ******************************************************************************/
#include <iostream>
#include <vector>
#include <algorithm>
#include <string>
#include <map>
#include <sstream>
#include <fstream>
#include <cstring>
#include <iomanip>
#include <unistd.h>
#include <unordered_map>
#include <unordered_set>
#include <signal.h>
#include <chrono>
#include <pthread.h>
#include <sched.h>
#include <sys/types.h>
#include "pcerrc.h"
#include "pmu.h"
#include "symbol.h"

static std::map<unsigned, double> numaTotalDDRC;                  //numa Id --> average ddrc bandwidth
static std::unordered_map<unsigned, unsigned*> numaToCpuCore;     //numa Id --> cpu core ids
static std::unordered_map<unsigned, unsigned> numaToCpuNumber;    //numa Id --> length of cpu cores
static std::vector<int> pidBoundCpus;                             //bounded cpus of designated pid
static unsigned numaNum = 0;                                      //number of NUMAs

const int FLOAT_PRECISION = 2;
const int TIME_UNIT_TRANS = 1000;

uint64_t topNum = 0;
uint64_t duration = 0;
uint64_t period = 0;

// 存储线程ID和对应的LLC miss率
struct ProcessMissRatio {
    int pid;
    unsigned llc_miss;
    unsigned llc_cache;
    double ratio;

    ProcessMissRatio(int p, unsigned miss, unsigned cache) : pid(p), llc_miss(miss), llc_cache(cache) {
        ratio = (cache != 0) ? static_cast<double>(miss) / cache : 0.0;
    }

    bool operator<(const ProcessMissRatio& other) const {
        return ratio > other.ratio; // 降序排列
    }
};

// 获取系统在线CPU核心数
int getSystemCpuCount() {
    return sysconf(_SC_NPROCESSORS_ONLN);
}

// 获取偶数核心列表
std::vector<int> getEvenCores() {
    std::vector<int> evenCores;
    int cpuCount = getSystemCpuCount();
    for (int i = 0; i < cpuCount; i += 2) {
        evenCores.push_back(i);
    }
    return evenCores;
}

// 获取奇数核心列表
std::vector<int> getOddCores() {
    std::vector<int> oddCores;
    int cpuCount = getSystemCpuCount();
    for (int i = 1; i < cpuCount; i += 2) {
        oddCores.push_back(i);
    }
    return oddCores;
}

std::vector<int> findStressNgProcesses() {
    std::vector<int> stressNgPids;
    DIR *dir = opendir("/proc");
    if (!dir) {
        std::cerr << "Failed to open /proc directory" << std::endl;
        return stressNgPids;
    }

    struct dirent *entry;
    while ((entry = readdir(dir)) != nullptr) {
        // 检查目录名是否为数字（PID）
        if (entry->d_type == DT_DIR) {
            char *endptr;
            long pid = strtol(entry->d_name, &endptr, 10);
            if (*endptr == '\0') {  // 确保整个字符串都是数字
                // 读取进程名
                std::string cmdlinePath = "/proc/" + std::string(entry->d_name) + "/cmdline";
                std::ifstream cmdlineFile(cmdlinePath);
                if (cmdlineFile.is_open()) {
                    std::string cmdline;
                    std::getline(cmdlineFile, cmdline);
                    // 检查命令行是否包含"stress-ng"
                    if (cmdline.find("stress-ng") != std::string::npos) {
                        stressNgPids.push_back(pid);
                    }
                }
            }
        }
    }
    closedir(dir);
    
    std::cout << "找到 " << stressNgPids.size() << " 个stress-ng进程" << std::endl;
    for (const auto& pid : stressNgPids) {
        std::cout << "PID: " << pid << std::endl;
    }
    
    return stressNgPids;
}

std::vector<ProcessMissRatio> collectProcessLlcMissData(const std::vector<int>& processList) {
    std::vector<ProcessMissRatio> processes;
    
    for (const auto& pid : processList) {
        // 获取进程的CPU亲和性
        if (getCpuCore(pid) == -1) {
            continue;
        }

        char* evtList[2];
        evtList[0] = (char*)"r33";
        evtList[1] = (char*)"r32";
        PmuAttr attr = {0};
        attr.evtList = evtList;
        attr.numEvt = 2;
        attr.pidList = const_cast<int*>(&pid);
        attr.numPid = 1;
        attr.cpuList = pidBoundCpus.data();
        attr.numCpu = pidBoundCpus.size();

        int pd = PmuOpen(COUNTING, &attr);
        if (pd == -1) {
            std::cerr << "PmuOpen failed for PID: " << pid << std::endl;
            std::cerr << "Error msg: " << Perror() << std::endl;
            continue;
        }

        PmuEnable(pd);
        
        // 收集数据
        unsigned llc_miss = 0;
        unsigned llc_cache = 0;
        
        // 收集一段时间的数据
        sleep(duration);
        
        PmuData* pmuData = nullptr;
        int len = PmuRead(pd, &pmuData);
        if (len == -1) {
            std::cerr << "PmuRead failed for PID: " << pid << std::endl;
            std::cerr << "Error msg: " << Perror() << std::endl;
            PmuDisable(pd);
            PmuClose(pd);
            continue;
        }
        
        // 处理收集到的数据
        for (int i = 0; i < len; ++i) {
            if (strcmp(pmuData[i].evt, "r33") == 0) {
                llc_miss += pmuData[i].count;
            }
            if (strcmp(pmuData[i].evt, "r32") == 0) {
                llc_cache += pmuData[i].count;
            }
        }
        
        PmuDataFree(pmuData);
        PmuDisable(pd);
        PmuClose(pd);
        
        // 添加到进程列表
        processes.emplace_back(pid, llc_miss, llc_cache);
    }
    
    return processes;
}

// 将线程绑定到指定的CPU核心
bool bindThreadToCore(int tid, int coreId) {
    cpu_set_t cpuset;
    CPU_ZERO(&cpuset);
    CPU_SET(coreId, &cpuset);

    // 使用sched_setaffinity绑定线程到指定核心
    if (sched_setaffinity(tid, sizeof(cpu_set_t), &cpuset) == -1) {
        std::cerr << "Failed to bind thread " << tid << " to core " << coreId << ": " << strerror(errno) << std::endl;
        return false;
    }

    std::cout << "Successfully bound thread " << tid << " to core " << coreId << std::endl;
    return true;
}

// 使用taskset命令绑定线程到指定核心（备选方法）
bool bindThreadWithTaskset(int tid, int coreId) {
    std::string cmd = "taskset -pc " + std::to_string(coreId) + " " + std::to_string(tid);
    int result = system(cmd.c_str());
    if (result != 0) {
        std::cerr << "Failed to execute: " << cmd << std::endl;
        return false;
    }
    return true;
}

void totalDDRCBandwidth()
{
    PmuDeviceAttr devAttr[2];
    devAttr[0].metric = PMU_DDR_READ_BW;
    devAttr[1].metric = PMU_DDR_WRITE_BW;
    int pd = PmuDeviceOpen(devAttr, 2);
    PmuEnable(pd);
    sleep(1);
    PmuData *oriData = nullptr;
    int oriLen = PmuRead(pd, &oriData);
    PmuDeviceData *devData = nullptr;
    auto len = PmuGetDevMetric(oriData, oriLen, devAttr, 2, &devData);
    std::unordered_map<int, double> stats;
    for (int i = 0; i < len; ++i) {
        stats[devData[i].ddrNumaId] += devData[i].count / 1024 / 1024;
    }
    for (const auto &entry : stats) {
        int id = entry.first;
        double sum = entry.second;
        numaTotalDDRC[id] = sum;
    }
    numaNum = numaTotalDDRC.size();
    DevDataFree(devData);
    PmuDataFree(oriData);
    PmuDisable(pd);
}

// get numaId --> cpu core ids
void initNumaToCoreList()
{
    unsigned *coreList;
    for (unsigned i = 0; i < numaNum; ++i) {
        coreList = nullptr;
        int len = PmuGetNumaCore(i, &coreList);
        numaToCpuCore[i] = coreList;
        numaToCpuNumber[i] = len;
    }
}

// 将进程绑定到指定的CPU核心
bool bindProcessToCore(int pid, int coreId) {
    cpu_set_t cpuset;
    CPU_ZERO(&cpuset);
    CPU_SET(coreId, &cpuset);

    // 使用sched_setaffinity绑定进程到指定核心
    if (sched_setaffinity(pid, sizeof(cpu_set_t), &cpuset) == -1) {
        std::cerr << "Failed to bind process " << pid << " to core " << coreId << ": " << strerror(errno) << std::endl;
        return false;
    }

    std::cout << "Successfully bound process " << pid << " to core " << coreId << std::endl;
    return true;
}

// 根据LLC miss率对进程进行排序并绑定到偶数/奇数核心
void sortAndBindProcesses(std::vector<ProcessMissRatio>& processes) {
    if (processes.empty()) {
        std::cout << "没有找到有效的进程数据" << std::endl;
        return;
    }
    
    // 按LLC miss率降序排序
    std::sort(processes.begin(), processes.end());
    
    // 获取系统CPU核心信息
    std::vector<int> evenCores = getEvenCores();
    std::vector<int> oddCores = getOddCores();
    
    if (evenCores.empty() || oddCores.empty()) {
        std::cerr << "无法获取CPU核心信息" << std::endl;
        return;
    }
    
    // 计算前N/2的进程数
    size_t halfProcesses = processes.size() / 2;
    
    std::cout << std::string(100, '=') << std::endl;
    std::cout << "绑定进程到CPU核心" << std::endl;
    std::cout << std::string(100, '-') << std::endl;
    std::cout << std::setw(10) << "序号" << std::setw(20) << std::left << "进程ID"
              << std::setw(20) << "LLC Miss率" << std::setw(20) << "绑定核心" << std::endl;
    std::cout << std::string(100, '-') << std::endl;
    
    // 绑定LLC miss率最高的进程到偶数核心
    for (size_t i = 0; i < halfProcesses && i < evenCores.size(); ++i) {
        int coreId = evenCores[i % evenCores.size()];
        bindProcessToCore(processes[i].pid, coreId);
        
        std::cout << std::setw(10) << i << std::setw(20) << std::left << processes[i].pid
                  << std::setw(20) << GetL3CMissPercent(processes[i].llc_miss, processes[i].llc_cache) + "%"
                  << std::setw(20) << coreId << " (偶数核)" << std::endl;
    }
    
    // 绑定剩余进程到奇数核心
    for (size_t i = halfProcesses; i < processes.size() && (i - halfProcesses) < oddCores.size(); ++i) {
        int coreId = oddCores[(i - halfProcesses) % oddCores.size()];
        bindProcessToCore(processes[i].pid, coreId);
        
        std::cout << std::setw(10) << i << std::setw(20) << std::left << processes[i].pid
                  << std::setw(20) << GetL3CMissPercent(processes[i].llc_miss, processes[i].llc_cache) + "%"
                  << std::setw(20) << coreId << " (奇数核)" << std::endl;
    }
    
    std::cout << std::string(100, '_') << std::endl;
}

// parse the CPU core range in the format "0-255" or "0-3,5"
std::vector<int> parseCpuRange(const std::string &rangeStr)
{
    std::vector<int> cpus;
    std::stringstream ss(rangeStr);
    std::string part;

    while(getline(ss, part, ',')) {
        size_t hyphen_pos = part.find("-");
        if (hyphen_pos != std::string::npos) {
            int start = std::stoi(part.substr(0, hyphen_pos));
            int end = std::stoi(part.substr(hyphen_pos + 1));
            if (start > end) {
                std::cerr << "Invalid CPU range: " << part << std::endl;
            }
            for (int i = start; i <= end; ++i) {
                cpus.push_back(i);
            }
        } else {
            cpus.push_back(std::stoi(part));
        }
    }

    std::sort(cpus.begin(), cpus.end());
    cpus.erase(unique(cpus.begin(), cpus.end()), cpus.end());
    return cpus;
}

// get cpu core of pid from /proc/[pid]/stat
std::string getCpuAffinityList(int pid)
{
    std::string path = "/proc/" + std::to_string(pid) + "/status";
    std::ifstream in(path);
    if (!in.is_open()) {
        std::cerr << "Not found: " << path << std::endl;
        return "";
    }
    std::string line;
    const std::string targetKey = "Cpus_allowed_list:";
    while (getline(in, line)) {
        if (line.find(targetKey) == 0) {
            size_t pos = line.find("\t");
            if (pos == std::string::npos)
                pos = targetKey.length();
            return line.substr(pos + 1);
        }
    }
    return "";
}

int getCpuCore(int pid)
{
    try {
        std::string rangeStr = getCpuAffinityList(pid);
        if (rangeStr == "") {
            return -1;
        }
        pidBoundCpus = parseCpuRange(rangeStr);
    } catch (const std::exception &e) {
        std::cerr << "Error: " << e.what() << std::endl;
        return 1;
    }
    return 0;
}

bool hasCommonCpu(const unsigned *cpuArray, size_t arraySize, const std::vector<int> &cpuVector)
{
    if (cpuArray == nullptr || arraySize == 0 || cpuVector.empty()) {
        return false;
    }

    if (arraySize < cpuVector.size()) {
        std::unordered_set<unsigned> arraySet(cpuArray, cpuArray + arraySize);
        for (const auto &cpu : cpuVector) {
            if (arraySet.count(cpu) > 0) {
                return true;
            }
        }
    } else {
        std::unordered_set<unsigned> vecSet(cpuVector.begin(), cpuVector.end());
        for (size_t i = 0; i < arraySize; ++i) {
            if (vecSet.count(cpuArray[i]) > 0) {
                return true;
            }
        }
    }

    return false;
}

std::string GetL3CMissPercent(unsigned llc_miss, unsigned llc_cache)
{
    std::ostringstream oss;
    double ratio = llc_cache != 0 ? static_cast<double>(llc_miss) / llc_cache * 100.0 : 0.0;
    oss << std::fixed << std::setprecision(FLOAT_PRECISION) << ratio;
    return oss.str();
}

// 收集线程的LLC miss数据并进行绑核处理
std::vector<ThreadMissRatio> CollectAndBindThreads(const std::unordered_map<unsigned, std::pair<unsigned, unsigned>>& tidData)
{
    std::vector<ThreadMissRatio> threads;

    // 将数据转换为ThreadMissRatio对象并存入vector
    for (const auto& pair : tidData) {
        unsigned tid = pair.first;
        unsigned llc_miss = pair.second.first;
        unsigned llc_cache = pair.second.second;
        threads.emplace_back(tid, llc_miss, llc_cache);
    }

    // 按LLC miss率降序排序
    std::sort(threads.begin(), threads.end());

    // 获取系统CPU核心信息
    std::vector<int> evenCores = getEvenCores();
    std::vector<int> oddCores = getOddCores();

    // 计算前N/2的线程数
    size_t halfThreads = threads.size() / 2;

    std::cout << std::string(100, '=') << std::endl;
    std::cout << "绑定线程到CPU核心" << std::endl;
    std::cout << std::string(100, '-') << std::endl;
    std::cout << std::setw(10) << "序号" << std::setw(20) << std::left << "线程ID"
              << std::setw(20) << "LLC Miss率" << std::setw(20) << "绑定核心" << std::endl;
    std::cout << std::string(100, '-') << std::endl;

    // 绑定前N/2的线程到偶数核心
    for (size_t i = 0; i < halfThreads && i < evenCores.size(); ++i) {
        int coreId = evenCores[i % evenCores.size()];
        bindThreadToCore(threads[i].tid, coreId);

        std::cout << std::setw(10) << i << std::setw(20) << std::left << threads[i].tid
                  << std::setw(20) << GetL3CMissPercent(threads[i].llc_miss, threads[i].llc_cache) + "%"
                  << std::setw(20) << coreId << " (偶数核)" << std::endl;
    }

    // 绑定剩余线程到奇数核心
    for (size_t i = halfThreads; i < threads.size() && (i - halfThreads) < oddCores.size(); ++i) {
        int coreId = oddCores[(i - halfThreads) % oddCores.size()];
        bindThreadToCore(threads[i].tid, coreId);

        std::cout << std::setw(10) << i << std::setw(20) << std::left << threads[i].tid
                  << std::setw(20) << GetL3CMissPercent(threads[i].llc_miss, threads[i].llc_cache) + "%"
                  << std::setw(20) << coreId << " (奇数核)" << std::endl;
    }

    std::cout << std::string(100, '_') << std::endl;

    return threads;
}

void PrintHotSpotGraph(const std::unordered_map<unsigned, std::pair<unsigned, unsigned>> tidData)
{
    std::vector<std::pair<unsigned, std::pair<unsigned, unsigned>>> sortedVec(tidData.begin(), tidData.end());
    std::sort(sortedVec.begin(), sortedVec.end(), [](const auto& a, const auto& b) {
        double ratioA = (a.second.second == 0) ? 0.0 : static_cast<double>(a.second.first) / a.second.second;
        double ratioB = (b.second.second == 0) ? 0.0 : static_cast<double>(b.second.first) / b.second.second;
        return ratioA > ratioB;
    });

    std::cout << std::string(100, '=') << std::endl;
    std::cout << std::string(100, '-') << std::endl;
    std::cout << " " << std::setw(10) << " " << std::setw(20) << std::left << "Tid" << std::setw(20) << "llc_cache_miss"
        << std::setw(20) << "llc_cache" << std::setw(20) << "llc_cache_miss_ratio" << std::endl;
    std::cout << std::string(100, '-') << std::endl;

    size_t outputNum = std::min(topNum, tidData.size());
    for (int i = 0; i < outputNum; ++i) {
        std::cout << " " << std::setw(10) << i << std::setw(20) << std::left << sortedVec[i].first << std::setw(20)
         << sortedVec[i].second.first << std::setw(20) << sortedVec[i].second.second << std::setw(20)
         << GetL3CMissPercent(sortedVec[i].second.first, sortedVec[i].second.second) + "%" << std::endl;
    }

    std::cout << std::string(100, '_') << std::endl;
}

int GetPmuDataHotspot(PmuData* pmuData, int pmuDataLen, bool bindThreads = false)
{
    if (pmuData == nullptr || pmuDataLen == 0) {
        return SUCCESS;
    }

    std::unordered_map<unsigned, std::pair<unsigned, unsigned>> tidData; //tid --> (0x33, 0x32)
    for (int i = 0; i < pmuDataLen; ++i) {
        PmuData& data = pmuData[i];
        if (strcmp(data.evt, "r33") == 0) {
            tidData[data.tid].first += data.count;
        }
        if (strcmp(data.evt, "r32") == 0) {
            tidData[data.tid].second += data.count;
        }
    }

    PrintHotSpotGraph(tidData);

    // 如果需要绑定线程到CPU核心
    if (bindThreads) {
        CollectAndBindThreads(tidData);
    }

    return SUCCESS;
}

void collectL3CMissRatio(int pid, bool bindThreads = false) {
    char* evtList[2];
    evtList[0] = (char*)"r33";
    evtList[1] = (char*)"r32";
    PmuAttr attr = {0};
    attr.evtList = evtList;
    attr.numEvt = 2;
    attr.pidList = &pid;
    attr.numPid = 1;
    attr.cpuList = pidBoundCpus.data();
    attr.numCpu = pidBoundCpus.size();

    int pd = PmuOpen(COUNTING, &attr);
    if (pd == -1) {
        std::cerr << "PmuOpen failed" << std::endl;
        std::cerr << "error msg:" << Perror() << std::endl;
        return;
    }

    PmuEnable(pd);
    int collectTimes = duration * TIME_UNIT_TRANS / period;
    for (int i = 0; i < collectTimes; ++i) {
        usleep(period * TIME_UNIT_TRANS);
        PmuData* pmuData = nullptr;
        int len = PmuRead(pd, &pmuData);
        if (len == -1) {
            std::cerr << "error msg:" << Perror() << std::endl;
            return;
        }

        // 最后一次收集时执行线程绑定
        bool shouldBindThreads = bindThreads && (i == collectTimes - 1);
        GetPmuDataHotspot(pmuData, len, shouldBindThreads);

        PmuDataFree(pmuData);
    }
    PmuDisable(pd);
    PmuClose(pd);
    return;
}

// g++ -o llc_miss_bind_cpu llc_miss_bind_cpu.cpp -I ./output/include/ -L ./output/lib/ -lkperf -lsym -lpthread
// export LD_LIBRARY_PATH=/XXX/libkperf/output/lib/:$LD_LIBRARY_PATH
void print_usage() {
    std::cerr << "Usage: llc_miss_bind_cpu <threshold> <topNum> <duration> <period>\n";
    std::cerr << "--threshold : the collect threshold of total ddrc bandwidth, unit M/s\n";
    std::cerr << "--topNum : the top N thread of llc miss ratio collection\n";
    std::cerr << "--duration : the total collect time of llc_miss_ratio, unit s\n";
    std::cerr << "--period : the period of llc_miss_ratio collect, unit ms\n";
    std::cerr << " example: llc_miss_bind_cpu 100 10 10 1000\n";
}

int main(int argc, char** argv)
{
    if (argc < 4) {
        print_usage();
        return 0;
    }
    double threshold = 0.0;
    bool collectL3CMissFlag = false;

    try {
        threshold = std::stod(argv[1]);
        if (threshold <= 0) {
            throw std::invalid_argument("threshold must be a positive number.");
        }

        topNum = std::stod(argv[2]);
        if (topNum <= 0) {
            throw std::invalid_argument("TopNum must be a positive number.");
        }

        duration = std::stod(argv[3]);
        if (duration <= 0) {
            throw std::invalid_argument("Duration must be a positive number.");
        }

        period = std::stoi(argv[4]);
        if (period <= 0) {
            throw std::invalid_argument("Period must be a positive integer.");
        }
    } catch (const std::exception& e) {
        std::cerr << "Error parsing arguments: " << e.what() << "\n";
        print_usage();
        return EXIT_FAILURE;
    }

    // 查找所有stress-ng进程
    std::vector<int> stressNgPids = findStressNgProcesses();
    if (stressNgPids.empty()) {
        std::cerr << "未找到任何stress-ng进程" << std::endl;
        return EXIT_FAILURE;
    }
    
    totalDDRCBandwidth();
    initNumaToCoreList();
    
    // 检查NUMA带宽是否超过阈值
    for (const auto &data : numaTotalDDRC) {
        std::cout << "Numa ID: " << data.first << ", total bandwidth: " << data.second << "M/s";
        if (data.second > threshold) {
            std::cout << " --> 超过阈值，需要收集LLC miss率";
            collectL3CMissFlag = true;
        } else {
            std::cout << " --> 未超过阈值";
        }
        std::cout << std::endl;
    }

    if (collectL3CMissFlag) {
        // 收集所有stress-ng进程的LLC miss率
        std::vector<ProcessMissRatio> processes = collectProcessLlcMissData(stressNgPids);
        
        // 根据LLC miss率排序并绑定到偶数/奇数核心
        sortAndBindProcesses(processes);
    }

    return 0;
}